#!/usr/bin/env python

'''
[macloc - search for host - find the port on a switch that a host belongs to]

Copyright (c) 2009, Kamil Derynski
All rights reserved.

Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:

    * Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
    * Neither the name of the <ORGANIZATION> nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
'''

__author__ = "Kamil Derynski"
__version__ = "0.5.2"

#for compatibility python2.5
try:
    Set = set
except NameError:
    from sets import Set
import commands
import netsnmp
import sys
import os
import time
import signal
from multiprocessing import Pool


#   --- Config section ---

#if you like floating screens ;)
debug = False

#SWITCH LIST  (Define here yours switch hostname list)
snlist = ["p507", "p510", "p225", "p226", "p226b", "p11"]

if debug:
    print snlist

#name etheraddres (Define file, where is name-ether mapping file)
etherfile = "/etc/ethers"

if debug:
    print etherfile

#Community snmp name
snmpname = "public"

if debug:
    print snmpname

#found it ?
foundit = False

#-----------------------------------------------------------------

class DictLookup(dict):
    u"""a dictionary which can lookup value by key, or keys by value"""
    def __init__(self, items=[]):
        u"""items can be a list of pair_lists or a dictionary"""
        dict.__init__(self, items)

    def get_key(self, value):
        u"""find the key(s) as a list given a value"""
        return [item[0] for item in self.items() if item[1] == value]

    def get_value(self, key):
        u"""find the value given a key"""
        return self[key]

def find(host):
    u"""Find host in ethers file"""
    ethers = parseEthers(etherfile) #dict
    result = {}

    try:
        key = ethers.get_key(host)
        if not key:
            raise KeyError
        #if mac address was found in ethers file
        result["name"] = key[0]
        result["mac"] = host
    except KeyError:
        try:
            result["name"] = host
            result["mac"] = ethers.get_value(host)
        except KeyError:
            #not found in ethers file
            raise EtherNotFound(result)
    return result

def parseEthers(file):
    u"""parse ethers file"""
    f=open(file, 'r')
    rows = f.readlines()
    out = {}
    for val in rows:
        #print val
        if val[0] == "#":
            pass
        else:
            l = val.split()
            if not l:
                pass
            else:
                out[l[1]] = l[0]
    f.close()
    return DictLookup(out)

class EtherNotFound(Exception):
    u'''Exception throw if Key Error ocured '''
    def __init__(self,msg):
        self.msg = msg

def cmd(command):
    u'''Console command function'''
    result = {}
    output = commands.getstatusoutput(command)
    result["status"] = output[0]
    result["out"] = output[1]
    if (int(result["status"]) != 0):
        raise CmdError(result)
    return result

def countDuplicatesInList(dupedList):
    u'''Count duplicate elements in list'''
    ret = {}
    uniqueSet = Set(item for item in dupedList)
    for item in uniqueSet:
        ret[item] = dupedList.count(item)
    return ret

class CmdError(Exception):
    u'''CmdError exception class '''
    def __init__(self, msg):
        if debug:
            self.msg = msg

class MacError(Exception):
    u'''MacError exception class '''
    def __init__(self):
        pass

def macHexToDec(mac):
    u"""Convert macaddress from xx:xx:xx:xx:xx:xx to dd.dd.dd.dd.dd.dd"""
    lMac = str(mac).split(":")
    if len(lMac) != 6:
        raise MacError
    try:
        lMac = [int(val,16) for val in lMac]
        macc = ".".join(["%s" % k for k in lMac])
    except ValueError:
        raise MacError
    return macc

class Host():
    u'''Contains information about host'''
    name = ""
    macHex = ""
    macDec = ""
    ip = ""

    def __init__(self,arg):
        try:
            self.macDec = macHexToDec(arg)
        except MacError:
            try:
                try:
                    self.macHex = find(arg)["mac"]
                except KeyError:
                    self.macHex = self.getEthFromPing(arg)
                except EtherNotFound:
                    self.macHex = self.getEthFromPing(arg)
                self.macDec = macHexToDec(self.macHex)
            except CmdError, e:
                if debug:
                    print 'Error occured with msg:\n ', e.msg["out"]
                print 'Error occured with msg:\n ', e
                sys.exit(2)
    
    def getEthFromPing(self,arg):
        cmd("ping -qc1 -t1 " + arg)
        return cmd("arp -n " + arg)["out"].split(" ")[3]

class Switch():
    name = ""
    alive = False
    ports = {}

    def __init__(self,name):
        self.name = name
        if self.isAlive():
            self.macTableIter()

    def isAlive(self):
        try:
            status = cmd("ping -qc1 -t1 " + self.name)["status"]
            ret = True
        except CmdError:
            snlist.remove(self.name)
            ret = False
        return ret

    def macTableIter(self):
        oid = netsnmp.Varbind(".1.3.6.1.2.1.17.4.3.1.2")
        vars = netsnmp.VarList(oid)
        sess = netsnmp.Session(Version = 1,DestHost = self.name,Community=snmpname)
        res =sess.walk(vars)
        
        tmp_ports = []
        if debug:
            print "======= " + self.name + " ========"
            for var in vars:
                print "  ",var.tag, var.iid, "=", var.val, '(',var.type,')'
            print "=================================================="
        
        for var in vars:
            tmp_ports.append(var.val)

        self.ports = countDuplicatesInList(tmp_ports)

    def portLocate(self,mac):
        oid = ".1.3.6.1.2.1.17.4.3.1.2." + mac
        var = netsnmp.Varbind(oid)
        res = netsnmp.snmpget(var,Version = 1,DestHost = self.name,Community=snmpname)
        if res[0] in self.ports:
            if self.ports[res[0]] < 2:
                print self.name + ":"+ res[0]
                return True
        return False

def worker(switch,hostmac):
    try:
        sw = Switch(switch)
        result = sw.portLocate(hostmac)
        print result
        return result
    except KeyboardInterrupt:
        print "watek"

def found(f):
    print foundit
    if f:
        foundit = True

def runMain():
    try:
        try:
            if not snlist:
                print "Define your switch list in source first !"
                sys.exit(2)
        except NameError:
            print "Define your switch list in source first !"
            sys.exit(2)


        try:
            arg = sys.argv[1]
        except IndexError:
            print "usage: macloc.py host"
            sys.exit(1)
        
        host = Host(arg)
        if debug:
            print "Host mac (dec): " +  host.macDec
        print "Host mac (dec): " +  host.macDec
        

        pool = Pool(processes=len(snlist))
        
        A = []
        for val in snlist:
            pool.apply_async(worker, (val,host.macDec), callback=A.append)
        
        pool.close()
        pool.join()
        
        print A
        if not foundit:
            print "Not found !"

    except KeyboardInterrupt:
        print "terminator"
        #pool.terminate()

if __name__ == '__main__':
    runMain()
